from .swinunet import get_swinunet
from .unet import UNet
from .transunet import get_transunet
from .segformer import SegFormer

def get_model(args):
    if args.model =='swinunet':
        model = get_swinunet(img_size=args.train_crop_size[0],num_classes=args.num_classes,in_channels=args.in_channels)
    elif args.model == 'unet':
        model = UNet(in_channels=args.in_channels,num_classes=args.num_classes)
    elif args.model == 'transunet':
        model=get_transunet(image_size=args.train_crop_size,in_channels=args.in_channels,num_classes=args.num_classes)
    elif args.model == 'segformer':
        model=SegFormer(image_size=args.train_crop_size,in_channels=args.in_channels,num_classes=args.num_classes)

    else:
        raise NotImplementedError

    return model